import operator
from typing import TypedDict,Annotated,Sequence

from langchain_community.chat_models import ChatOllama
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import StateGraph,START,END

from rich.console import Console
console = Console()


# 多模型选择
model_llama = ChatOllama(model="llama3.1")
model_gemma = ChatOllama(model="gemma2")

model = {
    "llama": model_llama,
    "gemma": model_gemma,
}
class AgentSate(TypedDict):
    messages: Annotated[Sequence[BaseMessage],operator.add]


def _call_model(state,config):
    model_name = config["configurable"].get("model", "llama")
    messages = state["messages"]
    if "system_messages" in config["configurable"]:
        messages = [
            SystemMessage(content=config["configurable"]["system_messages"])
        ] + messages
    m = model[model_name]
    resp = m.invoke(state["messages"])
    return {"messages": [resp]}

# 构建图
workflow = StateGraph(AgentSate)

workflow.add_node("call_model", _call_model)
workflow.add_edge(START, "call_model")
workflow.add_edge("call_model", END)

app = workflow.compile()

# 配置
config = {"configurable": {"system_messages": "说中文回复"}}

res = app.invoke({"messages": [HumanMessage(content="你好")]}, config=config)
print(res)